from helpers import normalize, blur

def initialize_beliefs(grid):
    '''
    :param grid:
    :return: 根据传入的空间,计算每个cell的平均概率
    '''
    height = len(grid)
    width = len(grid[0])
    area = height * width
    belief_per_cell = 1.0 / area
    beliefs = []
    for i in range(height):
        row = []
        for j in range(width):
            row.append(belief_per_cell)
        beliefs.append(row)
    return beliefs

def sense(color, grid, beliefs, p_hit, p_miss):
    '''
    :param color: 观察到的颜色
    :param grid: 移动的空间
    :param beliefs: 在空间中某处的概率分布
    :param p_hit: 走对的概率
    :param p_miss:  走错的概率
    :return: 根据观察到的结果(此处值指观察到的颜色)来修正之前的的概率

    使用贝叶斯公式更新概率后记得进行归一化操作
    '''
    new_beliefs = []

    #
    # TODO - implement this in part 2
    #

    # 获得感应后的概率分布
    for i,rows in enumerate(grid):
        it = []
        for j,item in enumerate(rows):
            if item == color:
                it.append(beliefs[i][j] * p_hit)
            else:
                it.append(beliefs[i][j] * p_miss)
        new_beliefs.append(it)

    # 归一化操作

    # 计算概率总和
    s = 0
    for i in new_beliefs:
        s += sum(i)

    # 计算归一化后的概率分布
    s1 = []
    for i in new_beliefs:
        s2 = []
        for j in i:
            s2.append(j / s)
        s1.append(s2)
    return s1

def move(dy, dx, beliefs, blurring):
    height = len(beliefs)
    width = len(beliefs[0])
    new_G = [[0.0 for i in range(width)] for j in range(height)]
    for i, row in enumerate(beliefs):
        for j, cell in enumerate(row):
            new_i = (i + dy) % height
            new_j = (j + dx) % width
            # pdb.set_trace()
            new_G[int(new_i)][int(new_j)] = cell

    return blur(new_G, blurring)
